from flask import Flask, jsonify, request, send_from_directory, abort
from flask_cors import CORS
from werkzeug.utils import secure_filename
from niqe import calculate_niqe_from_file_stream
from PIL import Image
import torch
import torch.nn as nn
from torchvision import transforms
import os
from io import BytesIO
import base64
from skimage.metrics import structural_similarity as ssim
from concurrent.futures import ThreadPoolExecutor,ProcessPoolExecutor
import numpy as np


# 初始化 Flask 应用
app = Flask(__name__)
CORS(app)  # 允许所有来源的跨域请求

# 配置上传文件的相关设置
UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
MAX_CONTENT_LENGTH = 16 * 1024 * 1024  # 16MB

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = MAX_CONTENT_LENGTH

# 确保上传目录存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# 定义网络结构
class Network(nn.Module):
    def __init__(self, n_chan, chan_embed=48):
        super(Network, self).__init__()
        self.act = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        self.conv1 = nn.Conv2d(n_chan, chan_embed, 3, padding=1)
        self.conv2 = nn.Conv2d(chan_embed, chan_embed, 3, padding=1)
        self.conv3 = nn.Conv2d(chan_embed, n_chan, 1)

    def forward(self, x):
        x = self.act(self.conv1(x))
        x = self.act(self.conv2(x))
        x = self.conv3(x)
        return x

# MSE 计算函数
def mse(gt: torch.Tensor, pred: torch.Tensor) -> torch.Tensor:
    loss = nn.MSELoss()
    return loss(gt, pred)

# 加载模型
def load_model(weights_path, n_chan=3, device='cuda'):
    model = Network(n_chan)
    model.load_state_dict(torch.load(weights_path, map_location=device))
    model.to(device)
    model.eval()
    return model

# 初始化模型
MODEL_WEIGHTS_PATH = './ckpt/Denoise_model.pth'
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
model = load_model(MODEL_WEIGHTS_PATH, n_chan=3, device=DEVICE)

# 图像转换器
transform = transforms.ToTensor()

# 辅助函数：将 PIL 图像转换为 Base64
def pil_to_base64(img: Image.Image) -> str:
    buffered = BytesIO()
    img.save(buffered, format="PNG")
    img_bytes = buffered.getvalue()
    return base64.b64encode(img_bytes).decode('utf-8')

# 辅助函数：去噪处理
def denoise_image(noise_img: Image.Image, device: str = DEVICE) -> Image.Image:
    noi = transform(noise_img).unsqueeze(0).to(device)  # BCHW
    with torch.no_grad():
        pred = torch.clamp(noi - model(noi), 0, 1)
    denoi = (pred * 255).squeeze(0).permute(1, 2, 0).cpu().numpy().astype(np.uint8)
    return Image.fromarray(denoi)

# 检查文件是否允许
def allowed_file(filename: str) -> bool:
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

# 处理单个文件，计算 NIQE 分数
def process_file(args):
    filename, file_bytes = args
    try:
        if allowed_file(filename):
            niqe_score = calculate_niqe_from_file_stream(file_bytes)
            niqe_score = round(niqe_score, 3) if niqe_score else None
            return filename, niqe_score
        return filename, None
    except Exception as e:
        print(f"Error processing file {filename}: {e}")
        return filename, None

# 初始化全局线程池
executor = ThreadPoolExecutor(max_workers=6)

# 路由：去噪
@app.route('/denoise/<int:index>', methods=['POST'])
def denoise_route(index):
    try:
        image_keys = [f'image{i}' for i in range(1, 4)]
        files = [(key, request.files[key]) for key in image_keys if key in request.files]

        denoised_images_b64 = []

        for key, file in files:
            img = Image.open(BytesIO(file.read())).convert('RGB')
            denoised_img = denoise_image(img)
            denoised_b64 = pil_to_base64(denoised_img)
            denoised_images_b64.append(denoised_b64)

        return jsonify({'denoised_images': denoised_images_b64}), 200
    except Exception as e:
        print(f"Error in denoise_route: {e}")
        return jsonify({'error': str(e)}), 500


# 路由：上传图片并计算 NIQE 分数
@app.route('/upload/<int:index>', methods=['POST'])
def upload_images(index):
    try:
        image_keys = [f'image{i}' for i in range(1, 7)]
        files = []

        for key in image_keys:
            if key not in request.files:
                return jsonify({'error': f'Missing file: {key}'}), 400
            file = request.files[key]
            if file.filename == '':
                return jsonify({'error': f'No selected file for {key}'}), 400
            if not allowed_file(file.filename):
                return jsonify({'error': f'File type not allowed for {key}'}), 400
            filename = secure_filename(file.filename)
            file_bytes = file.read()
            files.append((filename, file_bytes))

        # 提交任务到线程池
        futures = [executor.submit(process_file, file) for file in files]
        results = [future.result() for future in futures]

        response_data = {filename: score if score is not None else 'Error processing file' 
                         for filename, score in results}

        return jsonify({'data': response_data, 'results': [score for _, score in results]}), 200
    except Exception as e:
        print(f"Error in upload_images: {e}")
        return jsonify({'error': str(e)}), 500

# 路由：访问上传的图片
@app.route('/uploads/<filename>')
def uploaded_file(filename):
    try:
        return send_from_directory(app.config['UPLOAD_FOLDER'], filename)
    except FileNotFoundError:
        abort(404)

# 运行 Flask 应用
if __name__ == "__main__":
    app.run(debug=True)
